import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import argparse
import logging
from datetime import datetime, timedelta
from factor_validation import FactorValidator

def get_future_returns(self, trade_date, n_days=5):
        """
        计算从trade_date开始未来n_days的行业收益率
        """
def setup_logging():
    """
    设置日志记录配置
    """
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'validation.log')),
            logging.StreamHandler()
        ]
    )

def validate_date(date_str):
    """
    验证日期格式是否正确
    """
    try:
        datetime.strptime(date_str, '%Y%m%d')
        return True
    except ValueError:
        return False

def main():
    # 设置日志
    setup_logging()
    logger = logging.getLogger(__name__)
    
    try:
        # 创建命令行参数解析器
        parser = argparse.ArgumentParser(description='因子有效性验证框架')
        parser.add_argument('--start_date', type=str, default='20230311', help='回测起始日期，格式：YYYYMMDD')
        parser.add_argument('--end_date', type=str, default='20230315', help='回测结束日期，格式：YYYYMMDD')
        parser.add_argument('--n_days', type=int, default=5, help='未来收益计算天数')
        parser.add_argument('--factors', type=str, nargs='+', 
                          default=['density', 'consecutive', 'historical', 'strength'],
                          help='需要验证的因子列表')
        args = parser.parse_args()
        
        # 验证日期格式
        if not validate_date(args.start_date) or not validate_date(args.end_date):
            logger.error("日期格式错误，请使用YYYYMMDD格式")
            return
            
        if args.start_date > args.end_date:
            logger.error("起始日期不能晚于结束日期")
            return
    
    except Exception as e:
        logger.error(f"参数解析失败: {str(e)}")
        return
        
    # 初始化因子验证器
    validator = FactorValidator()
    
    # 创建输出目录
    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'output')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 验证每个因子
    for factor_name in args.factors:
        logger.info(f"\n{'='*50}")
        logger.info(f"开始验证因子: {factor_name}")
        logger.info(f"{'='*50}")
        
        try:
            # 运行因子验证
            validator.validate_factor(factor_name, args.start_date, args.end_date, args.n_days)
            
            # 运行市场状态测试
            print("\n4. 市场状态测试")
            try:
                market_state_ic = validator.run_market_state_test(factor_name, args.start_date, args.end_date, args.n_days)
                if market_state_ic and not all(np.isnan(v) for v in market_state_ic.values()):
                    print("不同市场状态下的平均IC值:")
                    for state, ic in market_state_ic.items():
                        print(f"{state}市场: {ic:.4f}")
                
                # 绘制市场状态测试图
                plt.figure(figsize=(8, 5))
                states = list(market_state_ic.keys())
                ics = [market_state_ic[s] for s in states]
                plt.bar(range(len(states)), ics)
                plt.xticks(range(len(states)), states)
                plt.title(f'{factor_name} 在不同市场状态下的表现')
                plt.xlabel('市场状态')
                plt.ylabel('平均IC')
                plt.grid(True)
                plt.tight_layout()
                
                # 保存图表
                plt.savefig(os.path.join(output_dir, f'{factor_name}_market_state.png'))
                plt.close()
            except Exception as e:
                print(f"无法进行有效的市场状态测试: {str(e)}")
        except Exception as e:
            logger.error(f"市场状态测试失败: {str(e)}", exc_info=True)
            continue
            
        # 运行风险调整收益分析
        print("\n5. 风险调整收益分析")
        try:
            risk_results = validator.run_risk_adjusted_returns(factor_name, args.start_date, args.end_date, args.n_days)
            if risk_results and 'high_cumulative' in risk_results and len(risk_results['high_cumulative']) > 0:
                print(f"最大回撤: {risk_results['max_drawdown']:.4f}")
                print(f"夏普比率: {risk_results['sharpe_ratio']:.4f}")
                print(f"Alpha: {risk_results['alpha']:.4f}")
                print(f"Beta: {risk_results['beta']:.4f}")
            
            # 绘制高分组、低分组和市场的累积收益率对比图
            plt.figure(figsize=(10, 6))
            plt.plot(risk_results['high_cumulative'], label='高分组')
            plt.plot(risk_results['low_cumulative'], label='低分组')
            plt.plot(risk_results['market_cumulative'], label='市场')
            plt.title(f'{factor_name} 风险调整收益分析')
            plt.xlabel('交易日序号')
            plt.ylabel('累积收益率')
            plt.legend()
            plt.grid(True)
            plt.tight_layout()
            
            # 保存图表
            plt.savefig(os.path.join(output_dir, f'{factor_name}_risk_adjusted.png'))
            plt.close()
        except Exception as e:
            logger.error(f"风险调整收益分析失败: {str(e)}", exc_info=True)
            continue
            
        logger.info(f"\n{'='*50}")
        logger.info(f"因子 {factor_name} 验证完成")
        logger.info(f"{'='*50}\n")
    
    logger.info("\n所有因子验证完成，结果已保存至output目录")

if __name__ == "__main__":
    main()